#alignment module:
import torch
import torch.nn as nn
from torchvision import ops

from .uNet_SKFF_v2 import UNet_SKFF_v2
from torch.nn import functional as F
from .attention_net import CBAM


class Offset_gen(nn.Module):
    def __init__(self, in_channels=9, out_channels=18):
        super(Offset_gen, self).__init__()

        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.Conv2d(64, out_channels, kernel_size=3, padding=1),
        )
        self.sigmod=nn.Sigmoid()

    def forward(self,x):  #输入tensor列表
        out = self.conv(x)

        return torch.tanh(out[:,0:18,...])*2, self.sigmod(out[:,18:,...])  # improved part


class Offset_gen_t(nn.Module):
    def __init__(self, in_channels=6, out_channels=18):
        super(Offset_gen_t, self).__init__()

        self.Offset_gen = UNet_SKFF_v2(nf=in_channels)
        self.offset_mask = nn.Conv2d(in_channels, out_channels, 3, padding=1)
         
        self.act=nn.LeakyReLU(negative_slope=0.1, inplace=True)
        self.sigmod=nn.Sigmoid()

    def forward(self,x,flag):  #输入tensor列表
        out = self.act(self.Offset_gen(x))
        out = self.offset_mask(out)

        if flag == 'pam':
            return torch.tanh(out[:,0:18,...])*2, self.sigmod(out[:,18:,...])
        else:    
            return out[:,0:18,...], self.sigmod(out[:,18:,...])


class Alignment(nn.Module):

    def __init__(self,cfg):
        super(Alignment,self).__init__()
        self.cfg = cfg
        # lack one pre conv
        self.offset_gens = nn.ModuleList()
        self.dconvs = nn.ModuleList()
        self.n_deform_conv = 2
        self.act = nn.LeakyReLU(negative_slope=0.1, inplace=True)

        self.fix_err = cfg.fix_err
        self.fix_mode = cfg.fix_mode

        if cfg.fix_err:
            self.fix_err_rate = cfg.fix_err_rate
            self.fix_errors = nn.Sequential(
                nn.Conv2d(3+2, 64, 3, 1, 1),
                nn.LeakyReLU(negative_slope=0.1, inplace=True),
                CBAM(64),
                nn.Conv2d(64, 2, 3, 1, 1),
                # nn.Tanh()
                nn.LeakyReLU(negative_slope=0.1, inplace=True)
            )
            if self.fix_mode == 'cat':
                self.offset_gens.append(Offset_gen_t(3+2+2, 18+9))
                self.offset_gens.append(Offset_gen_t(64+2+2, 18+9))
            else:
                self.offset_gens.append(Offset_gen_t(3+2, 18+9))
                self.offset_gens.append(Offset_gen_t(64+2, 18+9))

        else:
            self.offset_gens.append(Offset_gen_t(3+2, 18+9))
            self.offset_gens.append(Offset_gen_t(64+2, 18+9))
        # TODO add feature extractor
        # TODO pay att to offset

        # self.offset_gens.append(Offset_gen(3+2,18+9))
        # self.offset_gens.append(Offset_gen(64+2,18+9))
        # self.offset_gens.append(Offset_gen(64+2,18+9))
        # self.offset_gens.append(Offset_gen(64+2,18+9))

        self.dconvs.append(ops.DeformConv2d(3, cfg.pam_inner_dim, kernel_size=3, padding=1, stride=1))
        self.dconvs.append(ops.DeformConv2d(cfg.pam_inner_dim, cfg.pam_out_dim, kernel_size=3, padding=1, stride=1))
        # self.dconvs.append(ops.DeformConv2d(cfg.pam_inner_dim, cfg.pam_inner_dim, kernel_size=3, padding=1, stride=1))
        # self.dconvs.append(ops.DeformConv2d(cfg.pam_inner_dim, cfg.pam_out_dim, kernel_size=3, padding=1, stride=1))
  
    def forward(self,x):
        E0 = x[:,-2:,...]
        Ft = x[:,:-2,...]
        if self.fix_err:
            fixed_err = self.fix_errors(torch.cat((Ft, E0), 1)) * self.fix_err_rate
            if self.fix_mode == 'add':
                E0 = E0 + fixed_err
        # E0 = E0 + fixed_err
        for j in range(self.n_deform_conv):
            if self.fix_err and self.fix_mode == 'cat':
                offset, mask = self.offset_gens[j](torch.cat((Ft, E0, fixed_err), 1),flag='pam')
            else:
                offset, mask = self.offset_gens[j](torch.cat((Ft, E0), 1),flag='pam')
            supp_feature = self.dconvs[j](Ft, offset, mask)
            Ft = self.act(supp_feature)
        return Ft